
## # design matrix
## x <- cbind(rep(1,4),1:4)
## # correlation matrix
## K <- cbind(c(1,0,.5,.5),c(0,1,.5,.5),c(.5,.5,1,.5),c(.5,.5,.5,1))
## # response value
## y <- 1:4
## # group variable
## group <- gl(5,4)
## group <- gl(5,4,labels=LETTERS[1:5])
## # data frame
## dat <- data.frame(x=rep(1:4,5),group=gl(5,4),y=rnorm(20))

loglikelihood <- function(ratio,dat,K,method="ML"){
    q <- nrow(K)
    ## FIXME: when X is not intercept only
    p <- 1
    group <- dat$group
    M <- nlevels(group)
    ## FIXME: when data is unbalanced?
    Z <- diag(q)
    delta1 <- chol(solve(K))
    Zs <- rbind(Z,ratio*delta1)
    QR <- qr(Zs)
    R11 <- qr.R(QR)

    R <- list()
    H <- list()
    ## FIXME: when data is unbalanced?
    ni <- nrow(dat)/M # row numbers of design matrix X
    ## FIXME: when X is not intercept: Xi goto for
    Xi <- matrix(rep(1,ni),ncol=1)
    Xsi <- rbind(Xi,matrix(0,q,ncol(Xi)))
    for(i in 1:M){
      ysi <- c(with(dat,subset(y,group==levels(group)[i])),rep(0,q))
      R[[i]] <- qr.qty(QR,Xsi)
      H[[i]] <- qr.qty(QR,ysi)
    }

    R00i <- list()
    H0i <- list()
    for(i in 1:M)
      R00i[[i]] <- R[[i]][-(1:q),,drop=F]
    for(i in 1:M)
      H0i[[i]] <- H[[i]][-(1:q)]
    # stack R00i and H0i
    R00i <- do.call(rbind,R00i)
    H0i <- unlist(H0i)
    QR <- qr(R00i)
    R00 <- qr.R(QR)
    H0i <- qr.qty(QR,H0i)
    H0 <- H0i[1:nrow(R00)]
    H1 <- H0i[-c(1:nrow(R00))]
    # profiled log-likelihood
    N <- nrow(dat)
    # FIXME: when data is unbalanced?
    if(method=="ML")
      l <- 0.5*N*(log(N)-log(2*pi)-1 - log(sum(H1^2))) + M*log(abs(ratio^q*prod(diag(delta1))/prod(diag(R11)))) else
    if(method=="REML")
      l <- 0.5*(N-p)*(log(N-p)-log(2*pi)-1 - log(sum(H1^2))) -log(abs(prod(diag(R00)))) + M*log(abs(ratio^q*prod(diag(delta1))/prod(diag(R11)))) else
    cat("no such method!")
    return(l)
}


lmm <- function(dat,K,interval=c(0.01,2),method="ML"){
    q <- nrow(K)
    group <- dat$group
    M <- nlevels(group)
    ## FIXME: when data is unbalanced?
    Z <- diag(q)
    delta1 <- chol(solve(K))
    opti <- optimize(loglikelihood,interval=interval,dat=dat,K=K,method=method,maximum=T)
    ratio <-opti$maximum
    Zs <- rbind(Z,ratio*delta1)
    QR <- qr(Zs)
    R11 <- qr.R(QR)
    
    R <- list()
    H <- list()
    ## FIXME: when data is unbalanced?
    ni <- nrow(dat)/M # row numbers of design matrix X
    for(i in 1:M){
      ysi <- c(with(dat,subset(y,group==levels(group)[i])),rep(0,q))
      ## FIXME: when X is not intercept
      Xi <- matrix(rep(1,ni),ncol=1)
      Xsi <- rbind(Xi,matrix(0,q,ncol(Xi)))
      R[[i]] <- qr.qty(QR,Xsi)
      H[[i]] <- qr.qty(QR,ysi)
    }

    R00i <- list()
    R10i <- list()
    H0i <- list()
    H1i <- list()
    for(i in 1:M){
      R10i[[i]] <- R[[i]][1:q,,drop=F]
      R00i[[i]] <- R[[i]][-(1:q),,drop=F]
    }
    for(i in 1:M){
      H1i[[i]] <- H[[i]][1:q]
      H0i[[i]] <- H[[i]][-(1:q)]
    }
    # stack R00i and H0i
    R00i <- do.call(rbind,R00i)
    H0i <- unlist(H0i)
    QR <- qr(R00i)
    R00 <- qr.R(QR)
    H0i <- qr.qty(QR,H0i)
    H0 <- H0i[1:nrow(R00)]
    H1 <- H0i[-c(1:nrow(R00))]
    # BLUE and BLUP
    N <- nrow(dat)
    betahat <- solve(R00)%*%H0
    sigmahat <- sqrt(sum(H1^2)/N)
    bhat <- list()
    R11.inv <- solve(R11)
    for(i in 1:M)
    bhat[[i]] <- R11.inv%*%(H1i[[i]] - R10i[[i]]%*%betahat)
    # profiled log-likelihood
     # FIXME: when data is unbalanced?
    #l <- 0.5*N*(log(N)-log(2*pi)-1) - N*log(sqrt(sum(H1^2))) + M*log(abs(ratio^q*prod(diag(delta1))/prod(diag(R11))))
    res <- list(betahat=betahat,sigmahat=sigmahat,sigmaghat=sigmahat/ratio,bhat=bhat,loglikelihood=opti$objective)
    class(res) <- "lmm"
    return(res)
}

print.lmm <- function(x){
 cat("Fixed effects:\n",x$betahat,"\nRandom effects:\n","\nSigmag\n",x$sigmaghat,"\nSigma\n",x$sigmahat,"\nLog-likelihood:",x$loglikelihood,"\n")
}


x <- lmm(dat=dat,K=K,method="REML")
x
xx <- lmm(dat=dat,K=K)
xx
xxx <- lmm1(dat=dat,K=K)
xxx


##### optimization
## Data generation I: Yi = beta0 + bi + ei 
library(MASS)
M <- 1000
beta0 <- 1
K <- cbind(c(1,0,.5,.5),c(0,1,.5,.5),c(.5,.5,1,.5),c(.5,.5,.5,1))
b <- mvrnorm(n=M, rep(0, 4), Sigma=K)
b <- as.vector(t(b))
e <- mvrnorm(n=M, rep(0, 4), Sigma=diag(1,4))
e <- as.vector(t(e))
y <- beta0+b+e
#cov(t(matrix(y,4)))
dat <- data.frame(y=y,group=gl(M,4))

## test
lmm(1,dat=dat,K=K)
optimize(lmm,c(0.2,3),dat=dat,K=K,maximum=T)
optimize(loglikelihood,c(0.2,3),dat=dat,K=K,maximum=T)
?optimize
# [1] 1.026296

## plot
lik <- c()
ratio <- seq(0,3,by=0.01)[-1]
for(i in 1:length(ratio))
  lik[i] <- lmm(ratio[i],dat=dat,K=K)
plot(ratio,lik)
ratio[which(lik==max(lik))]

## Data generation II: Yi = beta0 + bi + ei
library(MASS)
M <- 1000
beta0 <- 1
K <- cbind(c(1,0,.5,.5),c(0,1,.5,.5),c(.5,.5,1,.5),c(.5,.5,.5,1))
b <- mvrnorm(n=M, rep(0, 4), Sigma=2*K)
b <- as.vector(t(b))
e <- mvrnorm(n=M, rep(0, 4), Sigma=diag(1,4))
e <- as.vector(t(e))
y <- beta0+b+e
#cov(t(matrix(y,4)))
dat <- data.frame(y=y,group=gl(M,4))

## test
lmm(1,dat=dat,K=K)
optimize(lmm,c(0.2,3),dat=dat,K=K,maximum=T)
# [1] 0.7045578

## Data generation III: Yi = beta0 + bi + ei
M <- 400
beta0 <- 1
K <- matrix(c(1,0,0,0,.5,0,.25,.25,.25,.25,
              0,1,0,0,.5,0,.25,.25,.25,.25,
              0,0,1,0,0,.5,.25,.25,.25,.25,
              0,0,0,1,0,.5,.25,.25,.25,.25,
              0,0,0,0,1,0,.5,.5,.5,.5,
              0,0,0,0,0,1,.5,.5,.5,.5,
              0,0,0,0,0,0,1,.5,.5,.5,
              0,0,0,0,0,0,0,1,.5,.5,
              0,0,0,0,0,0,0,0,1,.5,
              0,0,0,0,0,0,0,0,0,1
              ),10)
K <- t(K)+K-diag(1,10)
b <-   mvrnorm(n=M, rep(0, 10), Sigma=K)
b <- as.vector(t(b))
e <-   mvrnorm(n=M, rep(0, 10), Sigma=diag(1,10))
e <- as.vector(t(e))
y <- beta0+b+e
#cov(t(matrix(y,10)))
dat <- data.frame(y=y,group=gl(M,10))

## test
lmm(1,dat=dat,K=K)
optimize(lmm,c(0.2,3),dat=dat,K=K,maximum=T)
# [1] 0.9948055

## plot
lik <- c()
ratio <- seq(0,3,by=0.01)[-1]
for(i in 1:length(ratio))
  lik[i] <- lmm(ratio[i],dat=dat,K=K)
plot(ratio,lik)
ratio[which(lik==max(lik))]



## Data generation IV: Yi = beta0 + bi + ei
M <- 400
beta0 <- 1
K <- matrix(c(1,0,0,0,.5,0,.25,.25,.25,.25,
              0,1,0,0,.5,0,.25,.25,.25,.25,
              0,0,1,0,0,.5,.25,.25,.25,.25,
              0,0,0,1,0,.5,.25,.25,.25,.25,
              0,0,0,0,1,0,.5,.5,.5,.5,
              0,0,0,0,0,1,.5,.5,.5,.5,
              0,0,0,0,0,0,1,.5,.5,.5,
              0,0,0,0,0,0,0,1,.5,.5,
              0,0,0,0,0,0,0,0,1,.5,
              0,0,0,0,0,0,0,0,0,1
              ),10)
K <- t(K)+K-diag(1,10)
b <-   mvrnorm(n=M, rep(0, 10), Sigma=3*K)
b <- as.vector(t(b))
e <-   mvrnorm(n=M, rep(0, 10), Sigma=diag(1,10))
e <- as.vector(t(e))
y <- beta0+b+e
dat <- data.frame(y=y,group=gl(M,10))

## test
lmm(1,dat=dat,K=K)
optimize(lmm,c(0.2,3),dat=dat,K=K,maximum=T)
# [1] 0.5900726


### decomposition of S matrix into nx(n-q) matrix
X <- cbind(rep(1,4),1:4)
S <- diag(4)-X%*%solve(t(X)%*%X)%*%t(X)
s <- eigen(S);
nminusp <- sum(round(s$values,2))
A <- s$vectors[,1:nminusp]
round(A%*%t(A),3)
round(t(A)%*%A,3)
K <- cbind(c(1,0,.5,.5),c(0,1,.5,.5),c(.5,.5,1,.5),c(.5,.5,.5,1))
Z <- diag(1,4)

AZKZA <- t(A)%*%K%*%A
azkza <- eigen(AZKZA)
lambda <- azkza$values
UR <- A%*%azkza$vectors
UR
y=rnorm(4)
eta <- t(UR)%*%y
lambda
ratio
sum(eta^2/(lambda+ratio))

loglikelihood1 <- function(ratio,dat,K,method="ML"){
  delta <- ratio^2
  q <- nrow(K)
  group <- dat$group
  M <- nlevels(group)
  ## FIXME: when data is unbalanced?
  # Z <- diag(q)
  # k <- eigen(Z%*%K%*%t(Z))
  k <- eigen(K)
  xi <- k$values
  # FIXME:when data is unbalanced?
  ni <- nrow(dat)/M
  Xi <- matrix(rep(1,ni),ncol=1)
  S <- diag(ni)-Xi%*%solve(t(Xi)%*%Xi)%*%t(Xi)
  s <- eigen(S);
  nminusp <- sum(round(s$values,2))
  A <- s$vectors[,1:nminusp]
  AZKZA <- t(A)%*%K%*%A
  azkza <- eigen(AZKZA)
  lambda <- azkza$values
  UR <- A%*%azkza$vectors
  logR <- list()
  # FIXME:when data is unbalanced?
  for(i in 1:M){
    yi <- with(dat,subset(y,group==levels(group)[i]))
    eta <- t(UR)%*%yi
    logR[[i]] <- log(sum(eta^2/(lambda+delta)))

  }
  if(method=="ML")
    l <- 0.5*( M*ni*(log(ni)-log(2*pi)-1) - ni*sum(unlist(logR)) - M*sum(log(xi+delta)) ) else
  if(method=="REML")
    l <- 0.5*( M*nminusp*(log(nminusp)-log(2*pi)-1) - nminusp*sum(unlist(logR)) - M*sum(log(lambda+delta)) )
  return(l)
}

lmm1 <- function(dat,K,interval=c(0.01,2),method="ML"){
  opti <- optimize(loglikelihood1,interval=interval,dat=dat,K=K,method=method,maximum=T)
  ratio <- opti$maximum
  delta <- ratio^2
  q <- nrow(K)
  group <- dat$group
  M <- nlevels(group)
  ## FIXME: when data is unbalanced?
  # Z <- diag(q)
  # k <- eigen(Z%*%K%*%t(Z))
  k <- eigen(K)
  xi <- k$values
  # FIXME:when data is unbalanced?
  ni <- nrow(dat)/M
  Xi <- matrix(rep(1,ni),ncol=1)
  S <- diag(ni)-Xi%*%solve(t(Xi)%*%Xi)%*%t(Xi)
  s <- eigen(S);
  nminusp <- sum(round(s$values,2))
  A <- s$vectors[,1:nminusp]
  AZKZA <- t(A)%*%K%*%A
  azkza <- eigen(AZKZA)
  lambda <- azkza$values
  UR <- A%*%azkza$vectors
  logR <- list()
  # FIXME:when data is unbalanced?
  for(i in 1:M){
    yi <- with(dat,subset(y,group==levels(group)[i]))
    eta <- t(UR)%*%yi
    logR[[i]] <- log(sum(eta^2/(lambda+delta)))
    res <- list(loglikelihood=opti$objective)
  }
  return(res)
}

xxx <- lmm1(dat=dat,K=K)
xxx
xxxx <- lmm1(dat=dat,K=K,method="REML")
xxxx

## Data generation: Yi = beta0 + bi + ei 
library(MASS)
M <- 100
beta0 <- 1
K <- cbind(c(1,0,.5,.5),c(0,1,.5,.5),c(.5,.5,1,.5),c(.5,.5,.5,1))
b <-   mvrnorm(n=M, rep(0, 4), Sigma=K)
b <- as.vector(t(b))
e <-   mvrnorm(n=M, rep(0, 4), Sigma=diag(1,4))
e <- as.vector(t(e))
y <- beta0+b+e
cov(t(matrix(y,4)))
dat <- data.frame(y=y,group=gl(M,4))

## test
lmm1(1,dat=dat,K=K)
lmm1(1,dat=dat,K=K,method="REML")
lmm(1,dat=dat,K=K)
optimize(lmm1,c(0.2,3),dat=dat,K=K,maximum=T)
optimize(lmm1,c(0.2,3),dat=dat,K=K,method="REML",maximum=T)
## plot
lik <- c()
ratio <- seq(0,3,by=0.01)[-1]
for(i in 1:length(ratio))
  lik[i] <- lmm1(ratio[i],dat=dat,K=K)
plot(ratio,lik)
ratio[which(lik==max(lik))]


lik <- c()
ratio <- seq(0,3,by=0.01)[-1]
for(i in 1:length(ratio))
  lik[i] <- lmm1(ratio[i],dat=dat,K=K,method="REML")
plot(ratio,lik)
ratio[which(lik==max(lik))]



## Data generation: Yi = beta0 + bi + ei
M <- 40
beta0 <- 1
K <- matrix(c(1,0,0,0,.5,0,.25,.25,.25,.25,
              0,1,0,0,.5,0,.25,.25,.25,.25,
              0,0,1,0,0,.5,.25,.25,.25,.25,
              0,0,0,1,0,.5,.25,.25,.25,.25,
              0,0,0,0,1,0,.5,.5,.5,.5,
              0,0,0,0,0,1,.5,.5,.5,.5,
              0,0,0,0,0,0,1,.5,.5,.5,
              0,0,0,0,0,0,0,1,.5,.5,
              0,0,0,0,0,0,0,0,1,.5,
              0,0,0,0,0,0,0,0,0,1
              ),10)
K <- t(K)+K-diag(1,10)
b <-   mvrnorm(n=M, rep(0, 10), Sigma=K)
b <- as.vector(t(b))
e <-   mvrnorm(n=M, rep(0, 10), Sigma=diag(1,10))
e <- as.vector(t(e))
y <- beta0+b+e
K
cov(t(matrix(y,10)))
dat <- data.frame(y=y,group=gl(M,10))

## test
lmm1(1,dat=dat,K=K)
lmm1(1,dat=dat,K=K,method="REML")
optimize(lmm1,c(0.2,3),dat=dat,K=K,maximum=T)
optimize(lmm1,c(0.2,3),dat=dat,K=K,method="REML",maximum=T)
## plot
lik <- c()
ratio <- seq(0,3,by=0.01)[-1]
for(i in 1:length(ratio))
  lik[i] <- lmm1(ratio[i],dat=dat,K=K)
plot(ratio,lik)
ratio[which(lik==max(lik))]


lik <- c()
ratio <- seq(0,3,by=0.01)[-1]
for(i in 1:length(ratio))
  lik[i] <- lmm1(ratio[i],dat=dat,K=K,method="REML")
plot(ratio,lik)
ratio[which(lik==max(lik))]


## comparison of the two method
system.time(
            lmm1(1,dat=dat,K=K,method="REML")
          )

system.time(
            lmm(1,dat=dat,K=K)
          )



## using whole data for calculation
loglikelihood2 <- function(ratio,dat,K,method="ML"){
  delta <- ratio*ratio
   ## FIXME: when data is unbalanced?
  # Z <- diag(q)
  # k <- eigen(Z%*%K%*%t(Z))
  k <- eigen(K)
  xi <- k$values
  # FIXME:when data is unbalanced?
  N <- nrow(dat)
  X <- matrix(rep(1,N),ncol=1)
  S <- diag(N)-X%*%solve(t(X)%*%X)%*%t(X)
  s <- eigen(S);
  # FIXME:when X is not intercept only
  # nminusp <- sum(round(s$values,2))
  nminusp <- N-1
  A <- s$vectors[,1:nminusp]
  AZKZA <- t(A)%*%K%*%A
  azkza <- eigen(AZKZA)
  lambda <- azkza$values
  UR <- A%*%azkza$vectors
   # FIXME:when data is unbalanced?
  eta <- t(UR)%*%dat$y
  logR <- log(sum(eta^2/(lambda+delta)))
  if(method=="ML")
    l <- 0.5*( N*(log(N)-log(2*pi)-1) - N*logR - sum(log(xi+delta)) ) else
  if(method=="REML")
    l <- 0.5*( nminusp*(log(nminusp)-log(2*pi)-1) - nminusp*logR - sum(log(lambda+delta)))
  return(l)
}



lmm2 <- function(dat,K,interval=c(0.01,2),method="ML"){
  opti <- optimize(loglikelihood2,interval=interval,dat=dat,K=K,method=method,maximum=T)
  ratio <- opti$maximum
  delta <- ratio*ratio
  k <- eigen(K)
  xi <- k$values
  # FIXME:when data is unbalanced?
  N <- nrow(dat)
  X <- matrix(rep(1,N),ncol=1)
  S <- diag(N)-X%*%solve(t(X)%*%X)%*%t(X)
  s <- eigen(S);
  # FIXME:when X is not intercept only
  # nminusp <- sum(round(s$values,2))
  nminusp <- N-1
  A <- s$vectors[,1:nminusp]
  AZKZA <- t(A)%*%K%*%A
  azkza <- eigen(AZKZA)
  lambda <- azkza$values
  UR <- A%*%azkza$vectors
  # FIXME:when data is unbalanced?
  eta <- t(UR)%*%dat$y
  R <- sum(eta^2/(lambda+delta))
  # BLUE and BLUP
  temp <- t(X)%*%solve(K)
  betahat <- solve(temp%*%X)%*%temp%*%dat$y
  if(method=="ML")
     sigmaghat <- sqrt(R/N) else
  if(method=="REML")
     sigmaghat <- sqrt(R/nminusp) else
  cat("no such method!")
  res <- list(betahat=betahat,sigmahat=sigmaghat*ratio,sigmaghat=sigmaghat,loglikelihood=opti$objective)
    class(res) <- "lmm"
 return(res)
}


lis <- list()
for(i in 1:100)
  lis[[i]] <- K
K1 <- as.matrix(do.call(bdiag,lis))
optimize(loglikelihood2,c(0.2,3),dat=dat,K=K1,maximum=T)

xxx <- lmm(dat=dat,K=K)
xxx
xxxx <- lmm2(dat=dat,K=K1)
xxxx

xxx <- lmm(dat=dat,K=K,method="REML")
xxx
xxxx <- lmm2(dat=dat,K=K1,method="REML")
xxxx
